import os
import sys
import zipfile
import re
from pathlib import Path
import pyzipper
import datetime

# 设置基准目录
script_dir = os.path.dirname(os.path.abspath(sys.argv[0]))
base_dir = os.path.abspath(os.path.join(script_dir, ''))

# 压缩密码
zip_password = b""

def get_filename_without_extension(file_path):
    return Path(file_path).stem


def zip_files_in_directory(source_dir):
    """检查目录中的文件并打包 ZIP"""
    folder_name= os.path.basename(source_dir)

    # ZIP 文件的名称为子目录名
    zip_filename = folder_name + datetime.datetime.now().strftime('%Y%m%d%H%M') + ".zip"

    zip_path = os.path.join(source_dir, zip_filename)

    # 判断是否提供了密码
    if zip_password and len(zip_password) > 0:
        # 如果有密码，则创建加密的 ZIP 文件
        with pyzipper.AESZipFile(zip_path, 'w', compression=pyzipper.ZIP_DEFLATED, encryption=pyzipper.WZ_AES) as zipf:
            zipf.setpassword(zip_password.encode())  # 设置密码
            for foldername, subfolders, filenames in os.walk(source_dir):
                for filename in filenames:
                    file_path = os.path.join(foldername, filename)
                    # 获取相对路径以保持目录结构
                    arcname = os.path.relpath(file_path, source_dir)

                    # 跳过第一级目录中的 .zip 文件
                    if foldername == source_dir and filename.lower().endswith('.zip'):
                        print(f"跳过文件: {file_path}")
                        continue

                    zipf.write(file_path, arcname)
    else:
        # 如果没有密码，则创建普通的 ZIP 文件
        with pyzipper.ZipFile(zip_path, 'w', compression=pyzipper.ZIP_DEFLATED) as zipf:

            for foldername, subfolders, filenames in os.walk(source_dir):
                for filename in filenames:
                    file_path = os.path.join(foldername, filename)
                    # 获取相对路径以保持目录结构
                    arcname = os.path.relpath(file_path, source_dir)

                    # 跳过第一级目录中的 .zip 文件
                    if foldername == source_dir and filename.lower().endswith('.zip'):
                        print(f"跳过文件: {file_path}")
                        continue

                    zipf.write(file_path, arcname)

    print(f"创建 ZIP 文件: {zip_filename}，使用密码保护")
    return zip_filename


def check_and_zip(folder, subfile) -> str:
    """检查目录中的文件并打包 ZIP"""
    basename = get_filename_without_extension(subfile)

    exe_name= os.path.basename(subfile)

    # ZIP 文件的名称为子目录名
    zip_filename = basename + ".zip"

    zip_path = os.path.join(folder, zip_filename)

    # 判断是否提供了密码
    if zip_password and len(zip_password) > 0:
        # 如果有密码，则创建加密的 ZIP 文件
        with pyzipper.AESZipFile(zip_path, 'w', compression=pyzipper.ZIP_DEFLATED, encryption=pyzipper.WZ_AES) as zipf:
            zipf.setpassword(zip_password.encode())  # 设置密码
            zipf.write(os.path.join(folder, subfile), exe_name)  # 添加文件
    else:
        # 如果没有密码，则创建普通的 ZIP 文件
        with pyzipper.ZipFile(zip_path, 'w', compression=pyzipper.ZIP_DEFLATED) as zipf:
            zipf.write(os.path.join(folder, subfile), exe_name)  # 添加文件

    print(f"创建 ZIP 文件: {zip_filename}")
    return zip_filename

if __name__ == "__main__":
    if len(sys.argv) > 1:
        if os.path.isdir(sys.argv[1]):
            current_directory = sys.argv[1]
            zip_files_in_directory(current_directory)
        else:
            check_and_zip(sys.argv[1])
    else:
        current_directory = os.getcwd()
        zip_files_in_directory(current_directory)

input('任意键退出')
